import os
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent

from my_deep_agents_from_scratch.file_tools import ls, read_file, write_file
from my_deep_agents_from_scratch.state import DeepAgentState
from utils import format_messages, show_prompt


def main():
    load_dotenv(os.path.join("..", ".env"), override=True)
    model = os.getenv("MODEL")
    base_url = os.getenv("BASE_URL")
    api_key = os.getenv("_API_KEY")

    # from my_deep_agents_from_scratch.prompts import (
    #     LS_DESCRIPTION,
    #     READ_FILE_DESCRIPTION,
    #     WRITE_FILE_DESCRIPTION,
    # )
    # show_prompt(LS_DESCRIPTION)

    # File usage instructions
    FILE_USAGE_INSTRUCTIONS = """You have access to a virtual file system to help you retain and save context.                                  

    ## Workflow Process                                                                                            
    1. **Orient**: Use ls() to see existing files before starting work                                             
    2. **Save**: Use write_file() to store the user's request so that we can keep it for later                     
    3. **Read**: Once you are satisfied with the collected sources, read the saved file and use it to ensure that you directly answer the user's question."""

    # Add mock research instructions
    SIMPLE_RESEARCH_INSTRUCTIONS = """IMPORTANT: Just make a single call to the web_search tool and use the result provided by the tool to answer the user's question."""

    # Full prompt
    INSTRUCTIONS = (
            FILE_USAGE_INSTRUCTIONS + "\n\n" + "=" * 80 + "\n\n" + SIMPLE_RESEARCH_INSTRUCTIONS
    )
    # show_prompt(INSTRUCTIONS)

    # Mock search result
    search_result = """模型上下文协议（MCP）是由Anthropic开发的一种标准协议，旨在实现AI模型与外部系统（如工具、数据库及其他服务）的无缝集成。
    它作为标准化的通信层，使AI模型能够以一致高效的方式访问并利用来自不同数据源的信息。
    本质上，MCP通过为数据交换提供统一语言，极大简化了将AI助手连接到外部服务的过程。"""

    # Mock search tool
    @tool(parse_docstring=True)
    def web_search(
            query: str,
    ):
        """Search the web for information on a specific topic.

        This tool performs web searches and returns relevant results
        for the given query. Use this when you need to gather information from
        the internet about any topic.

        Args:
            query: The search query string. Be specific and clear about what
                   information you're looking for.

        Returns:
            Search results from search engine.

        Example:
            web_search("machine learning applications in healthcare")
        """
        return search_result

    # Create agent using create_react_agent directly
    # model = ChatOpenAI(model=model, base_url=base_url, api_key=api_key)
    from langchain_deepseek.chat_models import ChatDeepSeek
    model = ChatDeepSeek(model=model, temperature=0, api_key=api_key, base_url=base_url)
    tools = [ls, read_file, write_file, web_search]

    # Create agent with system prompt
    agent = create_react_agent(
        model, tools, prompt=INSTRUCTIONS, state_schema=DeepAgentState
    )

    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "给我一份模型上下文协议（MCP）的简要概述。",
                }
            ],
            "files": {},
        }
    )
    format_messages(result["messages"])
    print(result["files"])


if __name__ == '__main__':
    main()